package com.hundsun.aitest.service.impl;

import com.alibaba.fastjson.JSONArray;
import com.hundsun.aitest.common.CaseTypeEnum;
import com.hundsun.aitest.model.ai.AutoCase;
import com.hundsun.aitest.model.ai.ChatRequest;
import com.hundsun.aitest.model.ai.Message;
import com.hundsun.aitest.service.AutoCaseParser;
import com.hundsun.aitest.service.BaoSongDBService;
import com.hundsun.aitest.service.CheckCaseService;
import com.hundsun.aitest.service.LargeModelService;
import com.hundsun.aitest.util.YamlOperationUtils;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.util.*;

@Slf4j
@Service
public class CheckCaseServiceImpl implements CheckCaseService {
    @Autowired
    LargeModelService largeModelService;

    @Autowired
    AutoCaseParser  autoCaseParser;

    @Autowired
    BaoSongDBService baoSongDBService;

    private static final int MAX_RETRIES = 2;


    @Override
    public AutoCase generateSingleCheckCase(String checkCode) {
        int retries = 0;
        AutoCase autoCase = null;

        while (autoCase == null && retries < MAX_RETRIES) {
            try {
                String prompt = buildPrompt(checkCode);
                ChatRequest chatRequest = new ChatRequest();
                chatRequest.setMessages(Arrays.asList(new Message("user", prompt)));
                String aiResponse = largeModelService.sendRequest(chatRequest);
                String jsonString = largeModelService.parseResponse(aiResponse);
                log.info("{}的查询AI的模型的结果:\nprompt:\n{},\naiResponse:\n{},\n结果内容:\n{}", checkCode, prompt, aiResponse,jsonString);
                autoCase = autoCaseParser.parse(jsonString, CaseTypeEnum.CHECK.getValue(), checkCode);
                if (autoCase != null) {
                    autoCase.setCaseType("2");
                }
            } catch (Exception e) {
                log.error("解析大模型返回结果失败，checkCode={}", checkCode, e);
            }
            retries++;
        }

        if (autoCase == null) {
            log.warn("Max retries reached without success for checkCode: {}", checkCode);
        }

        return autoCase;
    }

    @Override
    public List<AutoCase> generateBatchCheckCases(List<String> checkIds) {
        List<AutoCase> batchAutoCases = new ArrayList<>();

        for (String checkId : checkIds) {
            AutoCase autoCase = generateSingleCheckCase(checkId);
            if (autoCase != null) {
                batchAutoCases.add(autoCase);
            } else{
                log.error("解析用例失败，勾稽编号: {}", checkId);
            }
        }

        return batchAutoCases;
    }

    private String buildPrompt(String checkCode) {
        //参数根据勾稽编码构造
        String checkRuleDescription = baoSongDBService.getCheckName(checkCode);
        String checkRuleDbStruct = baoSongDBService.getTableMetaData(checkCode);
        JSONArray columnComments = baoSongDBService.getColumnComments(checkCode);
        String sqlString=baoSongDBService.getCheckRuleSql(checkCode);

        String yamlPath = "prompt/CheckCaseGenerate.yaml";
        Map<String, Object> inputDataOne = new HashMap<>();
        inputDataOne.put("checkRuleDescription", checkRuleDescription);
        inputDataOne.put("checkRuleDbStruct", checkRuleDbStruct);
        inputDataOne.put("sqlString", sqlString);
        inputDataOne.put("columnComments", columnComments.toString());
        List<String> prompts = YamlOperationUtils.getPromptQuestion(yamlPath, inputDataOne);
        if(prompts!=null&&prompts.size()>0){
            return prompts.get(0);
        }
        return "";

    }
}
